from enum import EnumMeta
from numpy import save
from pycocotools.coco import COCO
import os
import shutil
from tqdm import tqdm
# import skimage.io as io
import matplotlib.pyplot as plt
import cv2
import numpy as np
from PIL import Image, ImageDraw
 
# 需要设置的路径
_id = 'b2923080f40311ebb708b06ebf3058c8'
dataDir= '/data/zmyu/dataset/%s/'#%_id 
# savepath="/data/wangqian/dataset/wheel/%s/"#%_id
savepath = "./outs/wheel/%s/"
img_dir=savepath+'images/'
anno_dir=savepath+'annotations/'
datasets_list=['']
expand_ration = 0.1 # 外扩尺寸

#coco有80类，这里写要提取类的名字，以person为例 
classes_names = ['wheel', 'truck']
#包含所有类别的原coco数据集路径
'''
目录格式如下：
$COCO_PATH
----|annotations
----|train2017
----|val2017
----|test2017
'''

 
headstr = """\
<annotation>
    <folder>VOC</folder>
    <filename>%s</filename>
    <source>
        <database>My Database</database>
        <annotation>COCO</annotation>
        <image>flickr</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>company</name>
    </owner>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>%d</depth>
    </size>
    <segmented>0</segmented>
"""
objstr = """\
    <object>
        <name>%s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>%d</xmin>
            <ymin>%d</ymin>
            <xmax>%d</xmax>
            <ymax>%d</ymax>
        </bndbox>
    </object>
"""
 
tailstr = '''\
</annotation>
'''
 
# 检查目录是否存在，如果存在，先删除再创建，否则，直接创建
def mkr(path):
    if not os.path.exists(path):
        os.makedirs(path)  # 可以创建多级目录


def id2name(coco):
    classes=dict()
    for cls in coco.dataset['categories']:
        classes[cls['id']]=cls['name']
    return classes
 

def write_xml(anno_path, head, objs, tail):
    with open(anno_path, "w") as f:
        f.write(head)
        for obj in objs:
            f.write(objstr%(obj[0],obj[1],obj[2],obj[3],obj[4]))
        f.write(tail)
 
 
def save_annotations_and_imgs(coco, dataset, filename, objs, dataDir, anno_dir, idx):
    #将图片转为xml，例:COCO_train2017_000000196610.jpg-->COCO_train2017_000000196610.xml
    dst_anno_dir = os.path.join(anno_dir, 'annotations')
    mkr(dst_anno_dir)
    
    img_path=dataDir+dataset+'/'+filename
    dst_img_dir = os.path.join(anno_dir, 'images')
    mkr(dst_img_dir)
    filename = '%d_'%idx + filename
    dst_imgpath=dst_img_dir+ '/' + filename
    img = Image.open(img_path)
    img = img.crop(objs.pop(-1)[1:])
    img.save(dst_imgpath)
    img = np.array(img)
    #if (img.shape[2] == 1):
    #    print(filename + " not a RGB image")
     #   return
    # shutil.copy(img_path, dst_imgpath)
    anno_path=dst_anno_dir + '/' + filename[:-3]+'xml'
    head=headstr % (filename, img.shape[1], img.shape[0], img.shape[2])
    tail = tailstr
    write_xml(anno_path, head, objs, tail)
 

def showimg(coco,dataset,img,classes,cls_id,show=True, dataDir=None, savepath=None):
    I=Image.open('%s/%s/%s'%(dataDir,dataset,img['file_name']))
    Ic = I.copy()
    #通过id，得到annotation
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None)
    # print(annIds)
    anns = coco.loadAnns(annIds)
    # print(anns)
    # coco.showAnns(anns)
    objs = []
    for ann in anns:
        class_name=classes[ann['category_id']]
        if class_name in classes_names:
            if 'bbox' in ann:
                bbox=ann['bbox']
                xmin = max(int(bbox[0]), 0)
                ymin = max(int(bbox[1]), 0)
                xmax = min(int(bbox[2] + bbox[0]), I.size[0]-1)
                ymax = min(int(bbox[3] + bbox[1]), I.size[1]-1)
                obj = [class_name, xmin, ymin, xmax, ymax]
                objs.append(obj)
                draw = ImageDraw.Draw(I)
                draw.rectangle([xmin, ymin, xmax, ymax], outline=tuple([0, 255, 0]))
    if show:
        plt.figure()
        plt.axis('off')
        plt.imshow(I)
        plt.show()

    recs, grp = get_wheel_zone(objs, I)

    for i,( objs, whl) in enumerate(zip(recs, grp)):
        grp[i].append(objs)
        # objs, whl = get_wheel_zone(objs, I)
        if objs:
            It = Ic.copy()
            I = It.crop((objs[1], objs[2], objs[3], objs[4]))
        else:
            return []
        show = True
        if show:
            for obj in whl:
                xmin, ymin, xmax, ymax = obj[1], obj[2], obj[3], obj[4]
                draw = ImageDraw.Draw(I)
                draw.rectangle([xmin, ymin, xmax, ymax], outline=tuple([0, 255, 0]))
        mkr(savepath+'/recs/')
        if not os.path.exists('outs/%d_'%i + img['file_name']):
            I.save(savepath+'/recs/%d_'%i + img['file_name'])
    # whl.append(objs)
    return grp


def get_wheel_zone(objs, image):
    col, raw = image.size
    truck, wheel = [], []
    for obj in objs:
        if obj[0] == 'truck':
            truck.append(obj[1:])
        else:
            wheel.append(obj+[(obj[1]+obj[3])/2, (obj[2]+obj[4])/2])
    if len(truck) == 0:
        return [], []
    truck.sort(key=lambda x:x[0])
    truck = truck[0] # 得到最左边的汽车框坐标
    whl = []
    xc = int((truck[0]+truck[2])/2) #得到汽车框的质心x值

    wheel.sort(key=lambda x:x[-2])

    # 得到车轮质心在框质心右边的所有轮子
    for w in wheel:
        if (w[-2] > truck[0] and w[-2] < truck[2]) \
            and (w[-1] > truck[1] and w[-1] < truck[3]) \
            and (w[-2] > xc) \
            and xc < col/2 \
            and w[-2] > xc+(truck[2]-xc)/3:
            whl.append(w)
    # 如果只有一个轮子就不裁剪
    if len(whl) <= 1:
        return [], []
    whl = [w[:5] for w in whl]
    for i, w in enumerate(whl):
        wid = (w[3] - w[1])/2
        whl[i].append(wid)
    grp = []
    temp = []
    for i, w in enumerate(whl):
        if not temp:
            temp.append(w)
        t = w[3] + w[-1]
        if i == len(whl)-1:
            grp.append(temp)
            break
        if t > whl[i+1][1]:
            temp.append(whl[i+1])
        else:
            if len(temp) > 1:
                grp.append(temp)
            temp = []
    recs = []
    for g in grp: #外接矩形
        temp = ['rec', min(g, key=lambda x:x[1])[1], # xmin
                       min(g, key=lambda x:x[2])[2], # ymin
                       max(g, key=lambda x:x[3])[3], # xmax
                       max(g, key=lambda x:x[4])[4]] # ymax 
        recs.append(temp)       
    #外扩
    for i, temp in enumerate(recs):
        deltax, deltay = temp[3] - temp[1], temp[4] - temp[2]
        recs[i][1] = max(0, temp[1] - int(expand_ration*deltax))
        recs[i][2] = max(0, temp[2] - int(expand_ration*deltay))
        recs[i][3] = min(col-1, temp[3] + int(expand_ration*deltax))
        recs[i][4] = min(raw-1, temp[4] + int(expand_ration*deltay))        

    # 求在裁剪后的图上的坐标
    for j, g in enumerate(grp):
        for i, w in enumerate(g):
            grp[j][i][1] -= recs[j][1]
            grp[j][i][2] -= recs[j][2]
            grp[j][i][3] -= recs[j][1]
            grp[j][i][4] -= recs[j][2]     

           
    if len(grp) < 1:
        return [], []
    return recs, grp

    # 求最大外接矩形
    temp = ['rec', min(whl, key=lambda x:x[1])[1], # xmin
                   min(whl, key=lambda x:x[2])[2], # ymin
                   max(whl, key=lambda x:x[3])[3], # xmax
                   max(whl, key=lambda x:x[4])[4]] # ymax
    if temp[-2] < col/2:
        return [], []

    # 外扩0.2倍
    deltax, deltay = temp[3] - temp[1], temp[4] - temp[2]
    temp[1] = max(0, temp[1] - int(expand_ration*deltax))
    temp[2] = max(0, temp[2] - int(expand_ration*deltay))
    temp[3] = min(col-1, temp[3] + int(expand_ration*deltax))
    temp[4] = min(raw-1, temp[4] + int(expand_ration*deltay))
     
    for i, w in enumerate(whl):
        whl[i][1] -= temp[1]
        whl[i][2] -= temp[2]
        whl[i][3] -= temp[1]
        whl[i][4] -= temp[2]

    return temp, whl


def extract(classes_names, dataDir, savepath, annFile):
    for dataset in datasets_list:
        #./COCO/annotations/instances_train2017.json
        # annFile='{}/annotations/instances_{}.json'.format(dataDir,dataset)
        # annFile = '/data/zmyu/dataset/coco_json/%s.json'%_id

        #使用COCO API用来初始化注释数据
        coco = COCO(annFile)
    
        #获取COCO数据集中的所有类别
        classes = id2name(coco)
        print(classes)
        #[1, 2, 3, 4, 6, 8]
        classes_ids = coco.getCatIds(catNms=classes_names)
        print(classes_ids)
        not_exists = 0
        for cls in classes_names:
            #获取该类的id
            cls_id=coco.getCatIds(catNms=[cls])
            img_ids=coco.getImgIds(catIds=cls_id)
            print(cls,len(img_ids))
            # imgIds=img_ids[0:10]
            for imgId in tqdm(img_ids):
                try:
                    img = coco.loadImgs(imgId)[0]
                except:
                    not_exists += 1
                    print('%d image not exists ...'%imgId)
                filename = img['file_name']
                # print(filename)
                objss=showimg(coco, dataset, img, classes, classes_ids, show=False, savepath=savepath, dataDir=dataDir)
                if not objss:
                    continue
                for i, objs in enumerate(objss):
                    save_annotations_and_imgs(coco, dataset, filename, objs, dataDir, anno_dir=savepath, idx=i)
            break
        print('not exists: %d'%not_exists)


if __name__ == '__main__':
    import yaml
    from threading import Thread

    parl = 1

    ctx = []
    with open('/data/wangqian/code/hzdet/config/data.yaml', 'r') as f:
        ctxs = yaml.safe_load(f)['train']
    thrs = []
    for ctx in ctxs:
        _id = ctx['img_path'].split('/')[-1]
        annFile = ctx['ann_path']
        savepath =  "./outs/wheel/%s/" % _id

        if not parl:
            extract(classes_names=classes_names, dataDir=ctx['img_path'], savepath=savepath, annFile=annFile)
            exit(0)
        else:
            t = Thread(target=extract, args=[classes_names, ctx['img_path'], savepath, annFile, ])
            t.start()
            thrs.append(t)

    for t in thrs:
        t.join()